﻿#include "H264Parser.h"

#include <QDebug>
#include <QCoreApplication>
#include <QDateTime>

H264Parser::H264Parser(QObject *parent)
    : QObject(parent)
{
    // 查找解码器
    h264_codec_ = (AVCodec *)avcodec_find_decoder(AV_CODEC_ID_H264);
    h264_parser_ = av_parser_init(h264_codec_->id);
    h264_codec_ctx_ = avcodec_alloc_context3(h264_codec_);

    // 对于x264编码器  禁用Annex B
    // 编码时禁用NAL合并
    av_opt_set(h264_codec_ctx_->priv_data, "annexb", "1", 0);
    av_opt_set(h264_codec_ctx_->priv_data, "repeat_headers", "0", 0);

    // 2. 打开解码器
    avcodec_open2(h264_codec_ctx_, h264_codec_, NULL);

    // 将AVFrame转为Jpeg图片
    // 1. 查找JPEG编码器
    //    jpeg_codec_ = avcodec_find_encoder(AV_CODEC_ID_MJPEG);
    //    // 2. 创建编码器上下文
    //    jpeg_ctx_ = avcodec_alloc_context3(jpeg_codec_);
}

H264Parser::~H264Parser()
{
    av_parser_close(h264_parser_);
    avcodec_free_context(&h264_codec_ctx_);
}

void H264Parser::inputData(const QByteArray &input)
{
    h264Catche.append(input);
}

/**
 * av_parser_parse2 是 FFmpeg 库中用于解析原始码流的函数，主要用于将输入的压缩数据分割成独立的帧或包。
 * int av_parser_parse2(AVCodecParserContext *s,
                     AVCodecContext *avctx,
                     uint8_t **poutbuf, int *poutbuf_size,
                     const uint8_t *buf, int buf_size,
                     int64_t pts, int64_t dts,
                     int64_t pos);
 *  AVCodecParserContext*	解析器上下文
 *  AVCodecContext*	编解码器上下文
 *  uint8_t**	输出缓冲区指针
 *  int*	输出缓冲区大小
 *  const uint8_t*	输入数据缓冲区
 *  int	输入数据大小
 *  int64_t	输入数据的显示时间戳
 *  int64_t	输入数据的解码时间戳
 *  int64_t	输入数据在源文件中的位置
 *  返回消耗的输入字节数
 *  如果返回0，表示需要更多数据才能完成解析
 */
QList<AVPacket*> H264Parser::parserPacket()
{
    QList<AVPacket*> packets;
    const uint8_t* data = (const uint8_t*)h264Catche.data();
    const uint8_t* dataPtr = (const uint8_t*)h264Catche.data();
    int size = static_cast<int>(h264Catche.size());

    while(size > 0) {
        uint8_t* out_data = nullptr;
        int out_len = 0;

        int consumed = av_parser_parse2(h264_parser_, h264_codec_ctx_,
                                        &out_data, &out_len,
                                        data, size,
                                        AV_NOPTS_VALUE, AV_NOPTS_VALUE, AV_NOPTS_VALUE);

        data += consumed;
        size -= consumed;

        if(out_len > 0) {
            AVPacket* pkt = av_packet_alloc();
            av_packet_from_data(pkt, static_cast<uint8_t*>(av_malloc(out_len)), out_len);
            memcpy(pkt->data, out_data, out_len);
            packets.append(pkt);
        }
    }

    // 移除已处理的数据
    int consumedSize = data - dataPtr;
    h264Catche.remove(0, consumedSize);

    return packets;
}

//QList<AVPacket*> H264Parser::parserPacket()
//{
//    QList<AVPacket*> packets;
//    const uint8_t* data = (const uint8_t*)h264Catche.data();
//    int size = static_cast<int>(h264Catche.size());
//    int total_consumed = 0;

//    while(size - total_consumed > 0)
//    {
//        uint8_t* out_data = nullptr;
//        int out_len = 0;

//        // 解析数据包
//        int consumed = av_parser_parse2(
//                    h264_parser_, h264_codec_ctx_,
//                    &out_data, &out_len,
//                    data + total_consumed, size - total_consumed,
//                    AV_NOPTS_VALUE, AV_NOPTS_VALUE, AV_NOPTS_VALUE
//                    );

//        if (consumed < 0) {
//            qWarning() << "Parser error:" << consumed;
//            break;
//        }

//        total_consumed += consumed;

//        if (out_len > 0)
//        {
//            // 直接引用解析器输出的内存（避免拷贝）
//            AVPacket* pkt = av_packet_alloc();
//            if (!pkt) {
//                qCritical() << "Failed to allocate AVPacket";
//                break;
//            }

//            av_new_packet(pkt, out_len);  // 分配内部缓冲区
//            memcpy(pkt->data, out_data, out_len);  // 复制数据

//            // 可选：验证起始码
//            if (pkt->size >= 4 &&
//                    pkt->data[0] == 0 && pkt->data[1] == 0 &&
//                    (pkt->data[2] == 1 || (pkt->data[2] == 0 && pkt->data[3] == 1))) {
//                packets.append(pkt);
//            }
//            else
//            {
//                qWarning() << "Invalid NAL unit start code";
//                av_packet_free(&pkt);
//            }
//        }
//    }

//    // 安全移除已处理数据
//    if (total_consumed > 0) {
//        h264Catche.remove(0, total_consumed);
//    }

//    return packets;
//}

// 创建SEI NAL单元
uint8_t *H264Parser::create_sei_nal_unit(const char *sei_text, int *sei_size)
{
    int payload_size = strlen(sei_text);
    int payload_type = 5; // 用户自定义SEI payload类型

    // 计算所需大小 (起始码4 + NAL头1 + payload类型1 + payload大小1 + payload + 结束标记1)
    int total_size = 4 + 1 + 1 + (payload_size < 0xFF ? 1 : 3) + payload_size + 1;
    uint8_t* sei_nal = (uint8_t*)av_malloc(total_size);
    if (!sei_nal) return NULL;

    // 填充起始码
    sei_nal[0] = 0x00;
    sei_nal[1] = 0x00;
    sei_nal[2] = 0x00;
    sei_nal[3] = 0x01;

    // NAL头 (类型6 = SEI)
    sei_nal[4] = 0x06;

    // SEI payload
    uint8_t* p = sei_nal + 5;

    // payload类型
    *p++ = payload_type;

    // payload大小 (EBSP编码)
    if (payload_size < 0xFF) {
        *p++ = payload_size;
    } else {
        *p++ = 0xFF;
        *p++ = (payload_size >> 8) & 0xFF;
        *p++ = payload_size & 0xFF;
    }

    // payload数据
    memcpy(p, sei_text, payload_size);
    p += payload_size;

    // RBSP尾部标记
    *p++ = 0x80;

    *sei_size = p - sei_nal;
    return sei_nal;
}

int H264Parser::insert_sei_into_packet(AVPacket *pkt, const char *sei_text)
{
    if(!pkt || !sei_text)
        return -1;

    // 1. 创建SEI NAL单元
    int sei_size = 0;
    uint8_t* sei_nal = create_sei_nal_unit(sei_text, &sei_size);
    if (!sei_nal)
        return AVERROR(ENOMEM);

    // 2. 检查原始数据是否有起始码
    uint8_t* orig_data = pkt->data;
    int orig_size = pkt->size;

    // 3. 创建新数据缓冲区
    int new_size = sei_size + orig_size;
    uint8_t* new_data = (uint8_t*)av_malloc(new_size);
    if (!new_data) {
        av_free(sei_nal);
        return AVERROR(ENOMEM);
    }

    // 4. 组装新数据: SEI NALU + 原始数据
    memcpy(new_data, sei_nal, sei_size);
    memcpy(new_data + sei_size, orig_data, orig_size);

    // 5. 更新AVPacket
    if (pkt->data) {
        av_free(pkt->data);
    }

    pkt->data = new_data;
    pkt->size = new_size;

    av_free(sei_nal);
    return 0;
}

bool H264Parser::is_h264_keyframe(const uint8_t *data, int size)
{
    // H.264 NAL 单元类型：0x05 = IDR帧（关键帧）
    if (size > 4 && (data[4] & 0x1F) == 5) {
        return true;
    }
    return false;
}

int H264Parser::decodePacket2Frame(const AVPacket &pkt, int& imageWidth, int& imageHeight)
{
    AVFrame* frame = av_frame_alloc();
    if (!frame) {
        return -1;
    }

    // 发送压缩数据到解码器
    int ret = avcodec_send_packet(h264_codec_ctx_, &pkt);
    if (ret < 0) {
        return ret;
    }

    // 接收解码后的帧
    while (ret >= 0) {

        ret = avcodec_receive_frame(h264_codec_ctx_, frame);
        if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF)
        {
            break;
        }
        else if (ret < 0)
        {
            qDebug() <<"Error during decoding";
            av_frame_free(&frame);
            return ret;
        }

        imageWidth = frame->width;
        imageHeight = frame->height;

        //        qDebug() << "image width: " << frame->width << " image height: " << frame->height << "key frame : " << frame->key_frame
        //                 << QByteArray((char*)pkt.data,pkt.size).left(10).toHex(' ');

#if 0
        QImage image = decodeFrame2QImage(*frame);
        if(!image.isNull())
        {
            QString filePath = QCoreApplication::applicationDirPath() + "/" + QDateTime::currentDateTime().toString("yy-MM-dd_hh-mm-ss-zzz.jpg");
            image.save(filePath, "JPEG");
        }
#endif

        av_frame_unref(frame);
    }
    av_frame_free(&frame);
    return 0;
}

QImage H264Parser::decodeFrame2QImage(const AVFrame &frame)
{
    AVFrame* rgb_frame = av_frame_alloc();
    rgb_frame->format = AV_PIX_FMT_RGB24;
    rgb_frame->width = frame.width;
    rgb_frame->height = frame.height;

    if (av_frame_get_buffer(rgb_frame, 32) < 0) {
        fprintf(stderr, "Could not allocate RGB frame\n");
        av_frame_free(&rgb_frame);
        return QImage();
    }

    // 创建转换上下文
    SwsContext* sws_ctx = sws_getContext(
                frame.width, frame.height, (AVPixelFormat)frame.format,
                frame.width, frame.height, AV_PIX_FMT_RGB24,
                SWS_BILINEAR, NULL, NULL, NULL);

    if (!sws_ctx) {
        fprintf(stderr, "Could not create scale context\n");
        av_frame_free(&rgb_frame);
        return QImage();
    }

    // 执行转换
    sws_scale(sws_ctx,
              frame.data, frame.linesize, 0, frame.height,
              rgb_frame->data, rgb_frame->linesize);

    // 创建QImage (需要复制数据，因为QImage不会管理FFmpeg的内存)
    QImage image(rgb_frame->width, rgb_frame->height, QImage::Format_RGB888);
    for (int y = 0; y < rgb_frame->height; y++) {
        memcpy(image.scanLine(y),
               rgb_frame->data[0] + y * rgb_frame->linesize[0],
                rgb_frame->width * 3);
    }

    // 清理
    sws_freeContext(sws_ctx);
    av_frame_free(&rgb_frame);

    return image;
}
